from openai import OpenAI

# If using auth
# auth = "user:pass"
# client = OpenAI(api_key="114514", base_url=f"http://{auth}@localhost:5000/v1")

# Not using auth
api_base_url = "http://localhost:5000/v1"
client = OpenAI(api_key="114514", base_url=api_base_url) # api_key随便填无所谓，base_url填写服务器IP和端口并添加/v1

input_text = """先程から何度も込められているため息に、確実に実体験からくる言葉だと分かって周としては胸が痛い。彼女にとって好意も悪意も向けられるのを慣れすぎている事が透けて見えて、やるせなさに唇を嚙んでしまう。
 彼女の今までの交友関係は真昼が天使様として振る舞ってきた軌跡ではあるのだが、それが全て快いものではないのだと改めて突き付けられていた。
 「可愛いものだと勉強の面倒見てもらえるから、評判がいい女と仲良くなって自分の評判を上げる、周りから除け者にされないように、とかそういうもの。悪いとまあ、その、アクセサリーというか戦利品？ として欲している殿方とか私に振られた男子を拾うために仲良くしてる振りをする方とか……まあ色々居た訳です」
  ややぐったりとげんなりを合わせたような声音は本当に苦労してきた事が窺えて、思わず労るように頭を撫でてしまう。
   真昼が思い出しただけで心労が嵩んでいそうな声と表情をしていたので周としてはお疲れ様という気持ちでいっぱいだった。""" # 要翻译的日文

query = "将下面的日文文本翻译成中文：" + input_text

gpt_dict = [
    {
        "src": "周",
        "dst": "周",
        "info": "名字，男孩",
    },
    {
        "src": "真昼",
        "dst": "真昼",
        "info": "名字，女孩",
    },
]
gpt_dict_text_list = []
for gpt in gpt_dict:
    src = gpt['src']
    dst = gpt['dst']
    info = gpt['info'] if "info" in gpt.keys() else None
    if info:
        single = f"{src}->{dst} #{info}"
    else:
        single = f"{src}->{dst}"
    gpt_dict_text_list.append(single)

gpt_dict_raw_text = "\n".join(gpt_dict_text_list)
# 0.10版本的user prompt，需要填写术语表，也可以留空。
query_v010 = "据以下术语表（可以为空）：\n" + gpt_dict_raw_text + "\n\n" + "将下面的日文文本根据上述术语表的对应关系和备注翻译成中文：" + input_text

extra_query = {
    'do_sample': True,
    'num_beams': 1,
    'repetition_penalty': 1.0,
}
models_list = client.models.list()
print(models_list)
# print(models_list.data[0].model_version)

# chat_completion = client.chat.completions.create(
for output in client.chat.completions.create(
    model="sukinishiro",
    ############# v0.9 prompt #############
    # messages=[
    #     {
    #         "role": "system",
    #         "content": "你是一个轻小说翻译模型，可以流畅通顺地以日本轻小说的风格将日文翻译成简体中文，并联系上下文正确使用人称代词，不擅自添加原文中没有的代词。"
    #     },
    #     {
    #         "role": "user",
    #         "content": f"{query}"
    #     }
    # ],
    ############# v0.10 prompt #############
    messages=[
        {
            "role": "system",
            "content": "你是一个轻小说翻译模型，可以流畅通顺地使用给定的术语表以日本轻小说的风格将日文翻译成简体中文，并联系上下文正确使用人称代词，注意不要混淆使役态和被动态的主语和宾语，不要擅自添加原文中没有的代词，也不要擅自增加或减少换行。"
        },
        {
            "role": "user",
            "content": f"{query_v010}"
        }
    ],
    temperature=0.1,
    top_p=0.3,
    max_tokens=512,
    frequency_penalty=0.0,
    seed=-1,
    extra_query=extra_query,
    stream=True,
    # stop=['\n\n']                     # NotImplement
    # n=1                               # NotImplement
    # logit_bias                        # won't support
    # response_format                   # won't support
    # tools                             # won't support
    # tool_choice                       # won't support
    # user                              # won't support
):
    # stream=True key response
    if output.choices[0].finish_reason:
        print("\nfinish reason is", output.choices[0].finish_reason)
    elif output.choices[0].delta.content:
        print(output.choices[0].delta.content, end="")

# stream=False key response

# translated_text = chat_completion.choices[0].message.content
# finish_reason = chat_completion.choices[0].finish_reason
# prompt_token_cnt = chat_completion.usage.prompt_tokens
# completion_token_cnt = chat_completion.usage.completion_tokens
# model = chat_completion.model

# print("\n".join([
#     f"The output is {translated_text}",
#     f"The inference stopped because of `{finish_reason}`.",
#     f"The output used {completion_token_cnt} tokens.",
#     f"The output was generated by the `{model}` model."
# ]))
